import os
import csv
import time
import datetime
import struct
import threading
import pandas as pd
import numpy as np
from pytdx.params import TDXParams
from pytdx.hq import TdxHq_API
from pytdx.exhq import TdxExHq_API
from chinese_calendar import is_holiday

pd.set_option('display.unicode.ambiguous_as_wide', True)
pd.set_option('display.unicode.east_asian_width', True)

# tdx 交易所代码
class Exchange():
    ZCE = 28
    DCE = 29
    SHFE = 30
    CFFEX = 47
    GFEX = 66
    ZCEO = 4
    DCEO = 5
    SHFEO = 6
    CFFEXO = 7
    GFEXO =67
    F=[ZCE, DCE, SHFE, CFFEX, GFEX]
    O=[ZCEO, DCEO, SHFEO, CFFEXO,  GFEXO]
    ALL=[ZCE, DCE, SHFE, CFFEX, GFEX, ZCEO, DCEO, SHFEO, CFFEXO, GFEXO]

    @staticmethod
    def str(exchage_id):
        dict_vars = vars(Exchange)
        for var_name,var_id in dict_vars.items():
            if var_id == exchage_id:
                if var_name[-1] == 'O':
                    var_name = var_name[0:-1]
                return var_name


def get_time():
    now = datetime.datetime.now()
    return now.strftime("%H:%M:%S")

# 获取合约列表
def get_contract_list(api, exchange_list):
    total_count = api.get_instrument_count()
    num = total_count/1000
    last_count = total_count%1000
    df = pd.DataFrame()
    for i in range(int(num)):
        data = api.get_instrument_info(i*1000,1000)
        df = df.append(data)
    data = api.get_instrument_info(num*1000,last_count)
    df = df.append(data)
    df = df[df.market.isin(exchange_list)]
    # df = df[df.code.str.contains('L9')]
    df.to_csv('contract_list.csv')
    print('生成合约列表完成')
    return df

#def 获取上个交易日
def get_last_trade_date(date,offset = 0):
    date -= datetime.timedelta(days=1)
    print(date)
    while is_holiday(date) or date.isoweekday() > 5:
        date -= datetime.timedelta(days=1)
    if offset != 0 :
        date += datetime.timedelta(days=offset)
    day = date.year
    day *= 100
    day += date.month
    day *= 100
    day += date.day
    return day

# 夜盘交易日期的真实日期
def get_real_date(day, time):
    if time >=900 and time <=1600:
        return day
    date = datetime.date(int(day/10000), int(day/100)%100, day%100)
    if time >=2100 and time <=2400:
        return get_last_trade_date(date)
    if time >= 0 and time <= 400:
        return get_last_trade_date(date, offset=1)

# k线数据转为c++二进制数据
def to_kline_binary(line):
    date = int(line['datetime'][0:10].replace('-', ''))
    byte = struct.pack('<iffffiif', date, line['open'],line['high'],line['low'],line['close'],line['position'],line['trade'], 0)
    return byte


# 分钟线数据转为c++二进制数据
def to_mline_binary(line):
    date = int(line['datetime'][0:10].replace('-', ''))
    time = int(line['datetime'][11:].replace(':', ''))
    date = get_real_date(date,time)
    time = int(time/100)*60+time%100
    print(date, line['position'],line['trade'])
    byte = struct.pack('<ihffffIIf', date,time, line['open'],line['high'],line['low'],line['close'],line['position'],line['trade'], 0)
    return byte


# 获取K线
def get_lines(api, market, stock_id, real_stock_id, line_type = TDXParams.KLINE_TYPE_DAILY, save_type='csv'):
    real_stock_id = real_stock_id.replace('L8', 'MAIN')
    real_stock_id = real_stock_id.replace('L9', 'INDEX')
    df = pd.DataFrame()
    for i in range(100000):
        for count in range(5): # 重试5次,防止获取失败
            data = api.get_instrument_bars(line_type, market, stock_id, i*600 , 600)
            if not data is None and len(data) > 0:
                break
        df = pd.DataFrame(data).append(df, ignore_index=True)
        if data is None or len(data) < 600:
            break
    if df.empty:
        print('获取数据为空')
        return df
    # df.index = pd.DatetimeIndex(df['datetime'])
    # df = df[:'2022-02-12 15:00']
    # df = df.loc[df.datetime <= '2022-05-19 15:00']

    path = 'data/' + dict_line_type[line_type] +'/' + Exchange.str(market)
    if not os.path.exists(path):
        os.makedirs(path)

    if save_type != 'csv':
        with open(path +'/'+ real_stock_id + '.'+dict_file_type[line_type],'ab+') as file:
            for index, line in df.iterrows():
                byte = to_kline_binary(line) if line_type == 4 else to_mline_binary(line)
                file.write(byte)
    else:
        df.to_csv(path +'/' + real_stock_id + '.csv')
    print('获取['+stock_id +']数据['+dict_line_type[line_type]+ ']成功!')
    return df

if __name__ == '__main__':
    # print(get_last_trade_date(datetime.date(2022,2,7),4))
    dict_line_type = {TDXParams.KLINE_TYPE_DAILY:'lday',TDXParams.KLINE_TYPE_1MIN:'minline',TDXParams.KLINE_TYPE_5MIN:'fzline'}
    dict_file_type = {TDXParams.KLINE_TYPE_DAILY:'day',TDXParams.KLINE_TYPE_1MIN:'lc1',TDXParams.KLINE_TYPE_5MIN:'lc5'}
    # 以下字段需要配置：
    exchange_list = Exchange.ALL
    exchange_list = [Exchange.CFFEX, Exchange.SHFE, Exchange.DCE, Exchange.ZCE, Exchange.GFEX]
    exchange_list = Exchange.F
    dict_line_type = {TDXParams.KLINE_TYPE_1MIN:'minline',TDXParams.KLINE_TYPE_5MIN:'fzline'}
    dict_file_type = {TDXParams.KLINE_TYPE_1MIN:'lc1',TDXParams.KLINE_TYPE_5MIN:'lc5'}
    # dict_line_type = {TDXParams.KLINE_TYPE_DAILY:'lday'}
    # dict_file_type = {TDXParams.KLINE_TYPE_DAILY:'day'}
    for type_value in dict_line_type.values():
        if not os.path.exists('data/'+type_value):
            os.makedirs('data/'+type_value)
    # 先获取市场的合约列表，根据合约获取合约的K线数据
    api = TdxHq_API()
    exapi = TdxExHq_API()
    if exapi.connect('39.108.101.114', 7727):
    # if exapi.connect('106.15.88.195', 7709):
        # contract_list = get_contract_list(exapi, exchange_list)
        # for index, contract in contract_list.iterrows():
            # for line_type in dict_line_type.keys():
                # market = contract.loc['market']
                # # 期权合约的name字段存储的是真正的合约代码
                # if market in Exchange.F:
                    # lines = get_lines(exapi, market, contract.loc['code'], contract.loc['code'], line_type,'binary')
                # if market in Exchange.O:
                    # lines = get_lines(exapi, market, contract.loc['code'], contract.loc['name'], line_type,'binary')
                # # lines = get_lines(exapi,contract.loc['market'], contract.loc['code'], line_type)
        # print("获取所有K线数据完成，合约数量[%d]" % len(contract_list))
        lines = get_lines(exapi, Exchange.ZCE, 'RSL9', 'RSMAIN', TDXParams.KLINE_TYPE_1MIN, 'binary')
        lines = get_lines(exapi, Exchange.DCE, 'LHL8', 'LHMAIN', TDXParams.KLINE_TYPE_5MIN, 'binary')
        lines = get_lines(exapi, Exchange.DCE, 'LHL8', 'LHMAIN', TDXParams.KLINE_TYPE_1MIN, 'binary')
        lines = get_lines(exapi, Exchange.DCE, 'PGL8', 'PGMAIN', TDXParams.KLINE_TYPE_5MIN, 'binary')
        lines = get_lines(exapi, Exchange.DCE, 'PGL8', 'PGMAIN', TDXParams.KLINE_TYPE_1MIN, 'binary')
        lines = get_lines(exapi, Exchange.DCE, 'RRL8', 'RRMAIN', TDXParams.KLINE_TYPE_5MIN, 'binary')
        lines = get_lines(exapi, Exchange.DCE, 'RRL8', 'RRMAIN', TDXParams.KLINE_TYPE_1MIN, 'binary')
        # lines = get_lines(exapi, Exchange.ZCE, 'OIL8', 'OIMAIN', TDXParams.KLINE_TYPE_5MIN, 'binary')
        # lines = get_lines(exapi, Exchange.ZCE, 'OIL8', 'OIMAIN', TDXParams.KLINE_TYPE_1MIN, 'binary')
        # lines = get_lines(exapi, Exchange.SHFE, 'ALL8', 'ALMAIN', TDXParams.KLINE_TYPE_5MIN, 'binary')
        # lines = get_lines(exapi, Exchange.SHFE, 'ALL8', 'ALMAIN', TDXParams.KLINE_TYPE_1MIN, 'binary')
        # lines = get_lines(exapi, 47, 'IC2210', 'IC2210',TDXParams.KLINE_TYPE_DAILY, 'binary')
        # mlines = get_lines(exapi, 47, 'IC2203', TDXParams.KLINE_TYPE_DAILY, 'binary')
        # mlines = get_lines(exapi, Exchange.GFEX, 'lc2404','lc2404', TDXParams.KLINE_TYPE_DAILY, 'binary')
        # mlines = get_lines(exapi, 47, 'IC2203', TDXParams.KLINE_TYPE_1MIN, 'binary')
        # mlines = get_lines(exapi, 30, 'ZNL9', TDXParams.KLINE_TYPE_1MIN)
        # mlines = get_lines(exapi, 4, 'CF7H11NK', TDXParams.KLINE_TYPE_5MIN,'binary')
        # mlines = get_lines(exapi,47, 'IC2203')
        # mlines = GetExMLine(exapi,28, 'AP2204')
        # mlines = GetExMLine(exapi,30, 'ZNL9')
        # print(mlines)
        exapi.disconnect()
    else:
        print("连接失败！")
